import argparse
import configparser
import csv
import os
from collections import defaultdict

import matplotlib.pyplot as plt
import numpy as np

##################################################################
# 决赛训练赛判题器
##################################################################

parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', type=str, default="C:/Users/99577/Desktop/huawei/third_round/dataset/data_raw")
parser.add_argument('--result_file', type=str, default="C:/Users/99577/Desktop/huawei/third_round/dataset/result/solution.txt")
parser.add_argument('--jpg_name', type=str, default="data_raw.jpg")
args = parser.parse_args()

DATA_DIR = args.data_dir
RESULT_FILE = args.result_file
JPG_FILE = f"{os.path.dirname(RESULT_FILE)}/{args.jpg_name}"
INHERIT_RATE = 0.05

with open(f"{DATA_DIR}/demand.csv", mode='rt', encoding='UTF-8') as _r:
    _reader = csv.reader(_r)
    _header = next(_reader)
    client_names = _header[2:]

with open(f"{DATA_DIR}/site_bandwidth.csv", mode='rt', encoding='UTF-8') as _r:
    _reader = csv.reader(_r)
    site_names = []
    site_bandwidth_dict = {}
    next(_reader)  # 去除表头
    for _site_name, _bandwidth in _reader:
        site_names.append(_site_name)
        site_bandwidth_dict[_site_name] = int(_bandwidth)

_config = configparser.RawConfigParser()
_config.read(f"{DATA_DIR}/config.ini", encoding='UTF-8')
qos_constraint = _config.getint('config', 'qos_constraint')
V = _config.getint('config', 'base_cost')
A = _config.getfloat('config', 'center_cost')
M = len(client_names)
N = len(site_names)

with open(f"{DATA_DIR}/qos.csv", mode='rt', encoding='UTF-8') as _r:
    _reader = csv.DictReader(_r)
    reachable_tuple_set = set()

    for _row in _reader:
        _site_name = _row['site_name']
        for _client_name in client_names:
            _qos = int(_row[_client_name])
            if _qos < qos_constraint:
                reachable_tuple_set.add((_client_name, _site_name))


def new_demand_iter():
    with open(f"{DATA_DIR}/demand.csv", mode='rt', encoding='UTF-8') as _r:
        _reader = csv.reader(_r)
        _header = next(_reader)
        client_names = _header[2:]

        _pre_mtime = None
        demand_tuple_dict = {}
        for _row in _reader:
            _new_mtime = _row[0]
            if (_pre_mtime != None) and (_pre_mtime != _new_mtime):
                yield demand_tuple_dict
                demand_tuple_dict = {}
            stream_id = _row[1]
            for client_name, deamnd in zip(client_names, _row[2:]):
                demand_tuple_dict[(client_name, stream_id)] = int(deamnd)
            _pre_mtime = _new_mtime
        yield demand_tuple_dict


def new_assign_iter():
    with open(RESULT_FILE, mode='rt', encoding='UTF-8') as _r:
        result_tuple_dict = {}
        _k = 0
        for _line in _r:
            _line = _line.strip('\n')
            client_name, assign_str = _line.split(':')
            if _k == M:
                yield result_tuple_dict
                result_tuple_dict = {}
                _k = 0
            assign_str = assign_str.strip('<>')
            for site_assign_str in assign_str.split('>,<'):
                stream_assign_list = site_assign_str.split(',')
                site_name = stream_assign_list[0]
                for stream_name in stream_assign_list[1:]:
                    result_tuple_dict[(client_name, stream_name)] = site_name
            _k += 1
        yield result_tuple_dict


edge_site_bw_list_dict = defaultdict(list)
center_site_bw_list = []

for t, (t_demand, t_assign) in enumerate(zip(new_demand_iter(), new_assign_iter())):
    if not t % 500:
        print(f"progress: t={t}")

    if t == 0:
        for site_name in site_names:
            edge_site_bw_list_dict[site_name].append(0)
    else:
        for site_name in site_names:
            inherit_bw = int(edge_site_bw_list_dict[site_name][-1] * INHERIT_RATE)
            edge_site_bw_list_dict[site_name].append(inherit_bw)

    site_stream_max_bw = defaultdict(int)

    for (client_name, stream_name), demand in t_demand.items():
        if not demand:
            continue
        assert (client_name, stream_name) in t_assign, f"需求未被满足: t={t}, client={client_name}, stream={stream_name}, demand={demand}"

        site_name = t_assign[(client_name, stream_name)]
        assert (client_name, site_name) in reachable_tuple_set, f"不可达的分配: t={t}, client={client_name}, stream={stream_name}, demand={demand}, site={site_name}"

        edge_site_bw_list_dict[site_name][-1] += demand
        site_stream_max_bw[(site_name, stream_name)] = max(site_stream_max_bw[(site_name, stream_name)], demand)

    center_site_bw_list.append(sum(site_stream_max_bw.values()))
    for site_name in site_names:
        assert edge_site_bw_list_dict[site_name][-1] <= site_bandwidth_dict[site_name],\
            f"超上限分配: t={t}, site={site_name}, assigned_bandwidth:{edge_site_bw_list_dict[site_name][-1]}, upper_bandwidth={site_bandwidth_dict[site_name]}"

T = t + 1

site_bw_sort_arr_dict = {site_name: np.sort(np.array(bw_list)) for site_name, bw_list in edge_site_bw_list_dict.items()}
site_bw_sort_arr_dict['Center'] = np.sort(np.array(center_site_bw_list))
site_bw_arg_sort_arr_dict = {site_name: np.argsort(np.array(bw_list)) for site_name, bw_list in edge_site_bw_list_dict.items()}
site_bw_arg_sort_arr_dict['Center'] = np.argsort(np.array(center_site_bw_list))

quantile_index_dict = {site_name: (int(np.ceil(0.95 * T) - 1)) for site_name in site_names}
quantile_index_dict['Center'] = (int(np.ceil(0.95 * T) - 1))

site_quantile_bw_dict = {site_name: arr[quantile_index_dict[site_name]] for site_name, arr in site_bw_sort_arr_dict.items()}
site_quantile_t_dict = {site_name: arr[quantile_index_dict[site_name]] for site_name, arr in site_bw_arg_sort_arr_dict.items()}


def _score(site_name, arr):
    quantile_bw = arr[quantile_index_dict[site_name]]
    if not arr[-1]:
        return 0
    elif quantile_bw <= V:
        return V
    else:
        return (int(quantile_bw - V)**2) / site_bandwidth_dict[site_name] + quantile_bw


site_score_dict = {site_name: _score(site_name, site_bw_sort_arr_dict[site_name]) for site_name in site_names}
edge_site_score = sum(site_score_dict.values())

site_score_dict['Center'] = site_quantile_bw_dict['Center'] * A
center_site_score = site_score_dict['Center']

total_score = edge_site_score + center_site_score
print(f"total score: {int(total_score)}, edge_site_score: {int(edge_site_score)}, center_site_score: {int(center_site_score)}")

sorted_site_name_by_score = sorted(site_names, key=site_score_dict.get, reverse=True)
sorted_site_name_by_score.insert(0, 'Center')
print("generate jpg ...")

site_bandwidth_dict['Center'] = -1

_col_num = 3
_row_num = int(np.ceil((N + 1)/3))
fig, axes_list = plt.subplots(_row_num, _col_num, figsize=(5*_col_num, 5 * _row_num))
for i, c in enumerate(sorted_site_name_by_score):
    axes = axes_list[i//_col_num][i % _col_num]
    axes.set_title(f"Site:{c}, Score:{int(site_score_dict[c])}, t:{site_quantile_t_dict[c]}"
                   f"\nQuantile:{site_quantile_bw_dict[c]}, BwUpper:{site_bandwidth_dict[c]}")
    axes.plot(site_bw_sort_arr_dict[c])
    axes.scatter([quantile_index_dict[c]], [site_bw_sort_arr_dict[c][quantile_index_dict[c]]], c='r', label='*')

fig.savefig(JPG_FILE, bbox_inches='tight')
print(f"save jpg to:{JPG_FILE}")
